from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
import base64
import os

class CryptoUtils:
    """加密工具类，提供对称加密和解密功能"""
    
    @staticmethod
    def encrypt(plaintext: str, key: str) -> str:
        """使用AES-256-CBC模式加密文本
        
        Args:
            plaintext: 需要加密的明文
            key: 密钥（将被填充或截断为32字节）
        
        Returns:
            str: 加密后的Base64编码字符串
        """
        # 确保密钥为32字节（AES-256需要32字节密钥）
        key_bytes = key.encode('utf-8')
        if len(key_bytes) > 32:
            key_bytes = key_bytes[:32]  # 截断
        else:
            key_bytes = key_bytes.ljust(32, b'\0')  # 填充
            
        # 生成随机IV（初始化向量）
        iv = os.urandom(16)
        
        # 创建加密器
        cipher = Cipher(
            algorithms.AES(key_bytes),
            modes.CBC(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()
        
        # 对明文进行填充（PKCS7）
        padder = padding.PKCS7(128).padder()  # AES块大小为128位
        padded_data = padder.update(plaintext.encode('utf-8')) + padder.finalize()
        
        # 加密
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()
        
        # 将IV和密文组合并进行Base64编码
        result = base64.b64encode(iv + ciphertext).decode('utf-8')
        return result
    
    @staticmethod
    def decrypt(ciphertext_b64: str, key: str) -> str:
        """使用AES-256-CBC模式解密密文
        
        Args:
            ciphertext_b64: 加密后的Base64编码字符串
            key: 密钥（将被填充或截断为32字节）
        
        Returns:
            str: 解密后的明文
        """
        try:
            # 确保密钥为32字节
            key_bytes = key.encode('utf-8')
            if len(key_bytes) > 32:
                key_bytes = key_bytes[:32]
            else:
                key_bytes = key_bytes.ljust(32, b'\0')
                
            # 解码Base64字符串
            encrypted_data = base64.b64decode(ciphertext_b64)
            
            # 提取IV（前16字节）
            iv = encrypted_data[:16]
            ciphertext = encrypted_data[16:]
            
            # 创建解密器
            cipher = Cipher(
                algorithms.AES(key_bytes),
                modes.CBC(iv),
                backend=default_backend()
            )
            decryptor = cipher.decryptor()
            
            # 解密
            padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
            
            # 去除填充
            unpadder = padding.PKCS7(128).unpadder()
            plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()
            
            return plaintext.decode('utf-8')
        except Exception as e:
            raise ValueError(f"解密失败: {str(e)}")

# 创建一个全局的加密工具实例
crypto_utils = CryptoUtils()